import os
import json
import subprocess
import argparse
from pathlib import Path

def run_ida(bin_file, result_file, ida_path, script_path):
    # 确保IDA路径和脚本路径是正确的
    if not os.path.exists(ida_path):
        raise FileNotFoundError(f"IDA路径不存在: {ida_path}")
    
    if not os.path.exists(script_path):
        raise FileNotFoundError(f"脚本路径不存在: {script_path}")
    
    if not os.path.exists(bin_file):
        raise FileNotFoundError(f"二进制文件不存在: {bin_file}")

    # 构造IDA命令，传入result_file作为参数
    ida_command = [
        os.path.join(ida_path, 'ida64.exe'),  # IDA的执行文件（Windows下为idaq.exe）
        '-B',  # 无图形界面模式
        "-S {} {}".format(script_path, result_file),  # 传递结果文件路径作为参数给脚本
        # result_file,
        bin_file,  # 要分析的二进制文件
    ]
    
#     path = [ida, "-A", "-S{} {} {}".format(self.script_path,log_folder,self.out_folder)]
#     path.append( bin_file)
# ret = run(path, env=os.environ.copy(), stdout=PIPE).returncode
    
    # 执行IDA命令
    try:
        print("正在运行IDA分析...")
        subprocess.run(ida_command, check=True)
        print(f"IDA分析完成，结果将保存到 {result_file}")
    except subprocess.CalledProcessError as e:
        print(f"执行IDA命令时出错: {e}")
    except Exception as e:
        print(f"发生错误: {e}")

def clear_mid_file(src_folder): # 删除IDA分析的中间文件
    for root, dirs, files in os.walk(src_folder):
        for f in files:
            if f.split('.')[-1] in ["id0", "id1", "nam", "til", "id2", "cfg", "i64", "idb"]:
                try:
                    os.remove(os.path.join(root, f))
                except:
                    pass

def run_juliet(bin_dir, result_root, ida_path, script_path):
    
    bin_dir = Path(bin_dir)
    result_root = Path(result_root)
    ida_path = Path(ida_path)
    script_path = Path(script_path)

    
    # 遍历 bin_dir 下的所有文件
    for entry in os.listdir(bin_dir):
        full_path = os.path.join(bin_dir, entry)
        if os.path.isfile(full_path):
            # 提取文件名，不含扩展名
            name, _ = os.path.splitext(entry)

            # 构造对应的结果文件路径
            result_dir = os.path.join(result_root, name)
            result_file = os.path.join(result_dir, 'result.json')

            # 检查 result.json 是否存在
            if not os.path.isfile(result_file):
                print(f"[警告] 未找到 {result_file}，跳过 {entry}")
                continue

            # 调用 run_ida
            print(f"处理 {entry} -> {result_file}")
            run_ida(full_path, result_file, ida_path, script_path)
            
    clear_mid_file(bin_dir)

if __name__ == '__main__':
    bin_dir = r"D:\Users\Administrator\Desktop\study\project\dataset\test\binary\good"
    result_root = r"D:\Users\Administrator\Desktop\study\project\dataset\test\json\good"
    ida_path = r'D:\Users\Administrator\Desktop\study\IDA_Pro_7.5_SP3\IDA_Pro_7.5_SP3'
    script_path = r'D:\Users\Administrator\Desktop\study\project\octopus\juliet_test.py'

    run_juliet(bin_dir, result_root, ida_path, script_path)
    